{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "238342c6",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from mpl_toolkits.mplot3d import Axes3D\n",
    "from PIL import Image\n",
    "\n",
    "class Drone:\n",
    "    def __init__(self, id, start_pos, target_pos):\n",
    "        self.id = id\n",
    "        self.start_pos = np.array(start_pos,dtype=\"float\")\n",
    "        self.target_pos = np.array(target_pos,dtype=\"float\")\n",
    "        self.current_pos = np.array(start_pos,dtype=\"float\")\n",
    "        self.path = [self.current_pos]\n",
    "\n",
    "    def move_towards_target(self, step_size):\n",
    "        direction = self.target_pos - self.current_pos\n",
    "        distance = np.linalg.norm(direction)\n",
    "        if distance > step_size:\n",
    "            self.current_pos += (direction / distance) * step_size\n",
    "        else:\n",
    "            self.current_pos = self.target_pos\n",
    "        self.path.append(self.current_pos)\n",
    "\n",
    "def ensure_min_distance(drones, min_distance):\n",
    "    for i in range(len(drones)):\n",
    "        for j in range(i + 1, len(drones)):\n",
    "            dist = np.linalg.norm(drones[i].current_pos - drones[j].current_pos)\n",
    "            if dist < min_distance:\n",
    "                direction = drones[i].current_pos - drones[j].current_pos\n",
    "                direction = direction / np.linalg.norm(direction)\n",
    "                drones[i].current_pos += direction * (min_distance - dist) / 2\n",
    "                drones[j].current_pos -= direction * (min_distance - dist) / 2\n",
    "\n",
    "def image_to_matrix(image_path, n):\n",
    "    image = Image.open(image_path).convert('L')  # Convert to grayscale\n",
    "    image = image.resize((n, n))  # Resize to n x n\n",
    "    matrix =4*(1-np.array(image) / 255.0)  # Normalize to [0, 1]\n",
    "    \n",
    "    matrix=np.array(matrix,dtype=\"int\")\n",
    "    print(np.mean(matrix))\n",
    "    return matrix\n",
    "\n",
    "def create_drones_from_matrix(matrix, ground_height=0):\n",
    "    drones = []\n",
    "    drone_id = 0\n",
    "    rows, cols = matrix.shape\n",
    "    for i in range(rows):\n",
    "        for j in range(cols):\n",
    "            num_drones = matrix[i, j]   # Scale to get a reasonable number of drones\n",
    "            for _ in range(num_drones):\n",
    "                start_pos = [np.random.rand()*rows, np.random.rand()*cols-25, ground_height]\n",
    "                target_pos = [j+0.5*int(_/2),0, cols-(i+0.5*int(_%2))]  # Target is directly above the start position\n",
    "                drones.append(Drone(drone_id, start_pos, target_pos))\n",
    "                drone_id += 1\n",
    "    return drones\n",
    "def change_drones(drones,matrix, ground_height=0):\n",
    "    drones_new=[]\n",
    "    drone_id = 0\n",
    "    rows, cols = matrix.shape\n",
    "    for i in range(rows):\n",
    "        for j in range(cols):\n",
    "            num_drones = int(matrix[i, j] *2)  # Scale to get a reasonable number of drones\n",
    "            for _ in range(num_drones):\n",
    "                if len(drones_new)>len(drones):\n",
    "                    start_pos=[np.random.rand()*rows, np.random.rand()*cols-25, ground_height]\n",
    "                else:\n",
    "                    start_pos = drones[drone_id].target_pos\n",
    "                target_pos = [j,0, cols-i]  # Target is directly above the start position\n",
    "                drones_new.append(Drone(drone_id, start_pos, target_pos))\n",
    "                drone_id += 1\n",
    "    if len(drones_new)<len(drones):\n",
    "        for i in range(len(drones_new),len(drones)):\n",
    "            start_pos = drones[i].target_pos\n",
    "            target_pos =[np.random.rand()*rows, np.random.rand()*cols-25, ground_height]\n",
    "            drones_new.append(Drone(drone_id, start_pos, target_pos))\n",
    "            drone_id += 1\n",
    "    return drones_new\n",
    "\n",
    "def simulate_drones(drones, step_size,min_distance,start_step=0,  max_steps=2000):\n",
    "    end_point=0\n",
    "    for step in range(start_step,max_steps):\n",
    "        fig = plt.figure(figsize=[10,10])\n",
    "        ax = fig.add_subplot(111, projection='3d')\n",
    "        ax.clear()\n",
    "        x=np.arange(0,50,1)\n",
    "        y=np.arange(-25,25,1)\n",
    "        X, Y=np.meshgrid(x, y)\n",
    "        ax.plot_surface(X,Y,Z=X*0-0.1,color=\"green\")\n",
    "        \n",
    "        all_at_target = True\n",
    "        for drone in drones:\n",
    "            if not np.array_equal(drone.current_pos, drone.target_pos):\n",
    "                all_at_target = False\n",
    "            \n",
    "            ax.scatter(drone.current_pos[0], drone.current_pos[1],drone.current_pos[2],color=\"black\", label=f'Drone {drone.id}')\n",
    "            #ax.plot([pos[0] for pos in drone.path], [pos[1] for pos in drone.path], [pos[2] for pos in drone.path], linewidth=10,color=\"black\")\n",
    "            drone.move_towards_target(step_size)\n",
    "        ensure_min_distance(drones, min_distance)\n",
    "        ax.set_xlim([0,50])\n",
    "        ax.set_ylim([-25,25])\n",
    "        ax.set_zlim([0,50])\n",
    "        ax.view_init(elev=e_angles[step%len(e_angles)], azim=a_angles[(step)%len(a_angles)])\n",
    "        ax.set_xlabel('X')\n",
    "        ax.set_ylabel('Y')\n",
    "        ax.set_zlabel('Z')\n",
    "        ax.set_title(f'Step {step}')\n",
    "        ax.axis(\"off\")\n",
    "        #plt.legend()\n",
    "        plt.savefig(SAVE_PATH+str(step)+\".jpg\")\n",
    "        plt.pause(0.01)\n",
    "        #np.array(target_pos,dtype=\"float\")\n",
    "        plt.show()\n",
    "        if all_at_target:\n",
    "            end_point+=1\n",
    "            if end_point==100:\n",
    "                return step\n",
    "    return max_steps\n",
    "# Example usage\n",
    "if __name__ == \"__main__\":\n",
    "    # Load image and convert to matrix\n",
    "    SAVE_PATH=r\"D:\\\\3_bodies\\\\live_game\\\\drone_\"\n",
    "    image_path = r\"D:\\\\3_bodies\\\\DLA\\\\金正恩.jpg\"  # Replace with your image path\n",
    "    image_path_1= r\"D:\\\\3_bodies\\\\DLA\\\\启动.jpg\"\n",
    "    e_angles =np.append(np.arange(0,15,0.1),np.arange(15,0,-0.1))\n",
    "    a_angles=np.arange(-180,180,1)\n",
    "    n = 50  # Size of the matrix\n",
    "    matrix = image_to_matrix(image_path, n)\n",
    "    plt.imshow(matrix,cmap=\"binary\")\n",
    "    plt.show()\n",
    "    # Create drones based on the matrix\n",
    "    drones = create_drones_from_matrix(matrix)\n",
    "    \n",
    "    # Simulate drones\n",
    "    end_step=simulate_drones(drones, step_size=0.1, min_distance=1)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f1fc36c6",
   "metadata": {},
   "outputs": [],
   "source": [
    "matrix = image_to_matrix(image_path_1, n)\n",
    "drones_1=change_drones(drones,matrix)\n",
    "end_step=simulate_drones(drones_1, step_size=0.1, min_distance=1,start_step=end_step)\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "53d334dc",
   "metadata": {},
   "outputs": [],
   "source": [
    "#imgs2video\n",
    "import os\n",
    "import cv2\n",
    "from PIL import Image\n",
    "fourcc = cv2.VideoWriter_fourcc('X', 'V', 'I', 'D') \n",
    "path=r\"D:/3_bodies/live_game/drone_\"#input(\"path?\")\n",
    "save_path=r\"D:/3_bodies/drone_3.mp4\"#input(\"save_path?\")\n",
    "post_video=cv2.VideoWriter(save_path, fourcc, 30,(720,720))\n",
    "len_img=633\n",
    "for i in range(len_img+1):\n",
    "    print(path+str(i)+\".jpg\")\n",
    "    img=cv2.imread(path+str(i)+\".jpg\")\n",
    "    post_video.write(img)\n",
    "    if i==len_img:\n",
    "        for t in range(100):\n",
    "            post_video.write(img)\n",
    "    del img\n",
    "    #if filename==\"0.jpg\":\n",
    "        #Image.fromarray(img).show()\n",
    "\n",
    "post_video.release()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0612a586",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
